from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import uvicorn
from typing import Dict, Set, Optional
import json
import time
import uuid

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        # 存储聊天室连接 {room_id: {connection_id: websocket}}
        self.room_connections: Dict[str, Dict[str, WebSocket]] = {}
        # 存储用户ID映射 {connection_id: user_info}
        self.user_sessions: Dict[str, dict] = {}
    
    async def connect(self, websocket: WebSocket, room_id: str, user_id: str, username: str):
        await websocket.accept()
        connection_id = str(uuid.uuid4())
        
        # 初始化聊天室
        if room_id not in self.room_connections:
            self.room_connections[room_id] = {}
        
        # 存储连接
        self.room_connections[room_id][connection_id] = websocket
        self.user_sessions[connection_id] = {
            "user_id": user_id,
            "username": username,
            "room_id": room_id
        }
        
        # 通知聊天室有新用户加入
        join_message = {
            "type": "system",
            "content": f"{username} 加入了聊天室",
            "time": time.strftime("%Y-%m-%d %H:%M:%S"),
            "room": room_id,
            "online_users": self.get_online_users(room_id)
        }
        await self.broadcast(json.dumps(join_message), room_id, exclude=[connection_id])
        
        return connection_id
    
    def disconnect(self, connection_id: str):
        if connection_id in self.user_sessions:
            user_info = self.user_sessions[connection_id]
            room_id = user_info["room_id"]
            
            # 从聊天室移除
            if room_id in self.room_connections and connection_id in self.room_connections[room_id]:
                del self.room_connections[room_id][connection_id]
                
                # 如果聊天室空了，清理
                if not self.room_connections[room_id]:
                    del self.room_connections[room_id]
            
            # 通知聊天室有用户离开
            if room_id in self.room_connections:
                leave_message = {
                    "type": "system",
                    "content": f"{user_info['username']} 离开了聊天室",
                    "time": time.strftime("%Y-%m-%d %H:%M:%S"),
                    "room": room_id,
                    "online_users": self.get_online_users(room_id)
                }
                asyncio.create_task(self.broadcast(json.dumps(leave_message), room_id))
            
            # 移除用户会话
            del self.user_sessions[connection_id]
    
    async def send_private_message(self, message: str, target_user_id: str, sender_info: dict):
        """发送私聊消息给特定用户"""
        target_connection = None
        target_username = None
        
        # 查找目标用户的连接
        for conn_id, user_info in self.user_sessions.items():
            if user_info["user_id"] == target_user_id:
                target_connection = self.room_connections[user_info["room_id"]].get(conn_id)
                target_username = user_info["username"]
                break
        
        if target_connection:
            try:
                private_msg = {
                    "type": "private",
                    "content": message,
                    "time": time.strftime("%Y-%m-%d %H:%M:%S"),
                    "from": sender_info,
                    "to": {"user_id": target_user_id, "username": target_username}
                }
                await target_connection.send_text(json.dumps(private_msg))
                
                # 同时给发送者返回确认消息
                confirmation = {
                    "type": "private_sent",
                    "content": f"私信已发送给 {target_username}",
                    "time": time.strftime("%Y-%m-%d %H:%M:%S"),
                    "to": {"user_id": target_user_id, "username": target_username}
                }
                return confirmation
            except Exception as e:
                print(f"发送私信时出错: {e}")
                return None
        else:
            return None
    
    async def broadcast(self, message: str, room_id: str, exclude: Optional[list] = None):
        """广播消息到聊天室，可以排除某些连接"""
        if room_id in self.room_connections:
            exclude = exclude or []
            for conn_id, connection in self.room_connections[room_id].items():
                if conn_id not in exclude:
                    try:
                        await connection.send_text(message)
                    except Exception as e:
                        print(f"广播消息时出错: {e}")
                        self.disconnect(conn_id)
    
    def get_online_users(self, room_id: str) -> list:
        """获取聊天室在线用户列表"""
        if room_id not in self.room_connections:
            return []
        
        return [
            {"user_id": self.user_sessions[conn_id]["user_id"], 
             "username": self.user_sessions[conn_id]["username"]}
            for conn_id in self.room_connections[room_id]
            if conn_id in self.user_sessions
        ]

manager = ConnectionManager()

@app.get("/")
async def get():
    return HTMLResponse("""
    <html>
        <head><title>WebSocket私聊示例</title></head>
        <body>
            <h1>WebSocket私聊示例服务端</h1>
            <p>支持多聊天室和私聊功能</p>
        </body>
    </html>
    """)

@app.websocket("/ws/{room_id}/{user_id}/{username}")
async def websocket_endpoint(
    websocket: WebSocket, 
    room_id: str, 
    user_id: str, 
    username: str
):
    # 建立连接并获取连接ID
    connection_id = await manager.connect(websocket, room_id, user_id, username)
    
    try:
        while True:
            data = await websocket.receive_text()
            
            try:
                json_data = json.loads(data)
                
                # 处理不同类型的消息
                if json_data.get("type") == "private":
                    # 私聊消息
                    target_user_id = json_data.get("to_user_id")
                    message_content = json_data.get("content", "")
                    
                    # 发送私聊
                    confirmation = await manager.send_private_message(
                        message_content, 
                        target_user_id,
                        {"user_id": user_id, "username": username}
                    )
                    
                    if confirmation:
                        await websocket.send_text(json.dumps(confirmation))
                
                else:
                    # 普通聊天室消息
                    message = {
                        "type": "chat",
                        "content": json_data.get("content", ""),
                        "time": time.strftime("%Y-%m-%d %H:%M:%S"),
                        "room": room_id,
                        "user": {"user_id": user_id, "username": username}
                    }
                    await manager.broadcast(json.dumps(message), room_id)
                    
            except json.JSONDecodeError:
                error_msg = {
                    "type": "error",
                    "content": "无效的JSON格式",
                    "time": time.strftime("%Y-%m-%d %H:%M:%S")
                }
                await websocket.send_text(json.dumps(error_msg))
                
    except WebSocketDisconnect:
        manager.disconnect(connection_id)
    except Exception as e:
        print(f"WebSocket错误: {e}")
        manager.disconnect(connection_id)

if __name__ == "__main__":
    import asyncio
    uvicorn.run(app, host="0.0.0.0", port=8000)